from pathlib import Path
import os
import math
import shutil


def read_data_ids(file_name):
    res = []
    with open(file_name, 'r') as file:
        for line in file:
            res.append(line.strip())
    return res

def copy_data_by_ids(root_dir, output_root_dir, data_ids):
    for data_id in data_ids:
        for folder in os.listdir(root_dir):
            if folder == data_id:
                dest_dir = output_root_dir / folder
                source_dir = os.path.join(root_dir, folder)
                shutil.copytree(source_dir, dest_dir.as_posix(), dirs_exist_ok=True)
                print(f'Copied: {source_dir} to {dest_dir.as_posix()}')


def copy_data(source_root_path, output_root_path, data_ids):
    output_dir = Path(output_root_path)
    output_dir.mkdir(parents=True, exist_ok=True)

    # 划分测试/训练集
    split_point = math.ceil(len(data_ids) * 103 / 150)

    train_data_ids = data_ids[:split_point]
    valid_data_ids = data_ids[split_point:]


    train_output_dir = output_dir / 'train'
    train_output_dir.mkdir(parents=True, exist_ok=True)

    train_source_root_path = os.path.join(source_root_path, 'train')
    valid_source_root_path = os.path.join(source_root_path, 'valid')

    # 训练数据
    copy_data_by_ids(train_source_root_path, train_output_dir, train_data_ids)
    copy_data_by_ids(valid_source_root_path, train_output_dir, train_data_ids)


    valid_output_dir = output_dir / 'valid'
    valid_output_dir.mkdir(parents=True, exist_ok=True)

    # 测试数据
    copy_data_by_ids(train_source_root_path, valid_output_dir, valid_data_ids)
    copy_data_by_ids(valid_source_root_path, valid_output_dir, valid_data_ids)


if __name__ == '__main__':
    daily_activity_list = read_data_ids('daily_activity.txt')
    animal_list = read_data_ids('animal.txt')
    artifact_landscape_list = read_data_ids('artifact_landscape.txt')

    root_path = '/mnt/DISK/xjk/undergraduate_thesis/dataset/DHF1K'
    output_root_path = '/mnt/DISK/xjk/undergraduate_thesis/dataset/CL'
    copy_data(root_path, output_root_path + '/daily_activity', daily_activity_list)
    copy_data(root_path, output_root_path + '/animal', animal_list)
    copy_data(root_path, output_root_path + '/artifact_landscape', artifact_landscape_list)













